#include "RayTraceRenderPass.h"
#include <SEngine/Render/RenderFeature/UtilRenderFeature.h>
#include <SERender/Renderer/Renderer.h>
#include <SEUi/SEUi.h>
#include <SERender/Resourses/Loader/ShaderLoader.h>
#include <SEngine/Render/RenderFeature/UBORenderFeature.h>
#include <SERender/Resourses/Loader/TextureLoader.h>
#include <SERender/Descriptors/FrameDescriptor.h>
#include <SEUi/Util/Util.h>

namespace SEngine
{
    RayTraceRenderPass::RayTraceRenderPass(Renderer & renderer)
        : RenderPass(renderer)
    {
        m_rayTraceMaterial = Manager::GetInstance().Get<Material>("RayTrace").value();
        m_billboardMaterial = Manager::GetInstance().Get<Material>("Billboard").value();
        m_fbo = std::make_unique<FrameBuffer>();
        TextureLoader<Texture2D> texLoader;
        texLoader.SetInternalFormat(RGBA16F);
        texLoader.SetFormat(RGBA);
        texLoader.SetType(FLOAT);
        texLoader.SetParameter(TEXTURE_MAG_FILTER, LINEAR);
        texLoader.SetParameter(TEXTURE_MIN_FILTER, LINEAR);
        texLoader.SetParameter(TEXTURE_WRAP_S, CLAMP_TO_EDGE);
        texLoader.SetParameter(TEXTURE_WRAP_T, CLAMP_TO_EDGE);    
        m_rayTraceRT[0] = texLoader.CreateTexture();
        m_rayTraceRT[1] = texLoader.CreateTexture();
    }

    void RayTraceRenderPass::DrawFrame()
    {
        auto & frame = m_renderer.AnyDescriptor::GetDescriptor<FrameDescriptor>()->get();
        auto & utilFeature = m_renderer.Descriptor<RenderFeature>::GetDescriptor<UtilRenderFeature>()->get();

        if (m_lastView != frame.camera.GetViewMatrix()) {
            this->Clear();
        }

        m_renderer.UseFrameBuffer(*m_fbo, [&]() {
            m_fbo->ClearAttachments();
            m_fbo->SwapAttachments(Attachment::COLOR_ATTACHMENT0, *m_rayTraceRT[m_frameCount % 2]);
            m_fbo->EnableRenderBuffer(true);
            m_fbo->Resize(frame.w, frame.h);
            m_fbo->Use();
            m_renderer.GetDrive()->Clear(COLOR_BUFFER_BIT | DEPTH_BUFFER_BIT);
            m_rayTraceMaterial.Value().Set("u_frameCount", m_frameCount);
            m_rayTraceMaterial.Value().Set("u_lastMap", Shader::Sampler{*m_rayTraceRT[(m_frameCount + 1) % 2]});
            utilFeature.DrawBillboard(m_rayTraceMaterial);
        });


        m_billboardMaterial.Value().Set("u_texture", Shader::Sampler{*m_rayTraceRT[m_frameCount % 2]});
        utilFeature.DrawBillboard(m_billboardMaterial);

        SignalUpdateUI.AddSlot([this]() {
            ImGui::PushID("GBufferRenderPass");
            if (ImGui::CollapsingHeader("GBufferRenderPass")) {
                if (ImGui::Button("Reload")) {
                    m_rayTraceMaterial.Value().GetShader().Reload([](auto & path, auto & name) {
                        ShaderLoader loader;
                        return loader.Load(path).CreateShader(EngineUniformBlockBinder());
                    });
                    this->Clear();
                }
            }
            m_rayTraceMaterial.Value().GetShader().Value().ForeachUniforms([this](const std::string & name, auto & type) {
                if constexpr (std::is_same_v<std::decay_t<decltype(type)>, float>) {
                    float v = m_rayTraceMaterial.Value().Get<float>(name).value_or(0.0f);
                    ImGui::DragFloat(name.c_str(), &v, 0.1f, 0.0f);
                    m_rayTraceMaterial.Value().Set(name, v);
                } else if constexpr (std::is_same_v<std::decay_t<decltype(type)>, glm::vec3>) {
                    glm::vec3 v = m_rayTraceMaterial.Value().Get<glm::vec3>(name).value_or(glm::vec3(0.0f));
                    ImGui::DragFloat3(name.c_str(), &v.x, 0.1f, 0.0f);
                    m_rayTraceMaterial.Value().Set(name, v);
                }
            });

            ImGui::PopID();
            return SlotState::Once;
        });
        m_frameCount++;
        m_lastView = frame.camera.GetViewMatrix();
    }

    void RayTraceRenderPass::Clear()
    {
        m_frameCount = 1;
        auto & frame = m_renderer.AnyDescriptor::GetDescriptor<FrameDescriptor>()->get();
        m_renderer.UseFrameBuffer(*m_fbo, [&]() {
            m_fbo->ClearAttachments();
            m_fbo->SwapAttachments(Attachment::COLOR_ATTACHMENT0, *m_rayTraceRT[0]);
            m_fbo->EnableRenderBuffer(true);
            m_fbo->Resize(frame.w, frame.h);
            m_fbo->Use();
            m_renderer.GetDrive()->Clear(COLOR_BUFFER_BIT | DEPTH_BUFFER_BIT);

            m_fbo->ClearAttachments();
            m_fbo->SwapAttachments(Attachment::COLOR_ATTACHMENT0, *m_rayTraceRT[1]);
            m_fbo->EnableRenderBuffer(true);
            m_fbo->Resize(frame.w, frame.h);
            m_fbo->Use();
            m_renderer.GetDrive()->Clear(COLOR_BUFFER_BIT | DEPTH_BUFFER_BIT);
        });
    }

} // namespace SEngine

